package main

import (
	"flag"
	"net"

	"github.com/sirupsen/logrus"
)

// 实现TCP层的端口转发服务
var (
	fromPort string
	toPort   string
)

func init() {
	fromPort = *flag.String("from", "8000", "need proxy port")
	toPort = *flag.String("to", "8080", "real server port")
	flag.Parse()
}

func main() {
	createTCPListener()
}

func createTCPListener() {
	listener, err := net.Listen("tcp", ":"+fromPort)
	if err != nil {
		logrus.Fatalf("server create tcp listener error: %v", err)
	}

	for {
		conn, err := listener.Accept()
		if err != nil {
			logrus.Errorf("listener create connect error%v", err)
			continue
		}

		go handleConnection(conn)
	}
}

func createTCPClient() net.Conn {
	clientConn, err := net.Dial("tcp", ":"+toPort)
	if err != nil {
		logrus.Fatalf("tcp client create error: %v", err)
	}

	return clientConn
}

func handleConnection(serverConn net.Conn) {
	clientConn := createTCPClient()

	go sendTarget2Source(serverConn, clientConn)
	sendSource2Target(serverConn, clientConn)
}

func sendSource2Target(serverConn, clientConn net.Conn) {
	var buf = [1024]byte{}

	for {
		num, err := serverConn.Read(buf[:])
		if err != nil {
			logrus.Errorf("tcp server connect read byte error: %v", err)
			serverConn.Close()
			clientConn.Close()
			break
		}

		_, err = clientConn.Write(buf[:num])
		if err != nil {
			logrus.Errorf("tcp client connect write byte error: %v", err)
			serverConn.Close()
			clientConn.Close()
			break
		}
	}
}

func sendTarget2Source(serverConn, clientConn net.Conn) {
	var buf = [1024]byte{}

	for {
		num, err := clientConn.Read(buf[:])
		if err != nil {
			logrus.Errorf("tcp client connect read byte error: %v", err)
			serverConn.Close()
			clientConn.Close()
			break
		}

		_, err = serverConn.Write(buf[:num])
		if err != nil {
			logrus.Errorf("tcp server connect write byte error: %v", err)
			serverConn.Close()
			clientConn.Close()
			break
		}
	}
}
